/*
    Copyright (c) Luo Yongping. THIS SOFTWARE COMES WITH NO WARRANTIES, 
    USE AT YOUR OWN RISKS!
*/

#include <iostream>
#include <libpmem.h>
#include <string>
#include <random>
#include <thread>
#include <cstring>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <omp.h>

#include "common.h"
#include "access.h"

using std::cout;
using std::endl;
using std::string;

double run_test(char * buf, int test_id, long size, long count, int thread_num) {
    drop_cache(); // clear the LLC cache
    
    auto start = seconds();
    // start all threads
    std::thread test_threads[thread_num];
    for(int i = 0; i < thread_num; i++) {
        if(test_id < 7) { // sequential
            long offset = i * count / thread_num * size;
            test_threads[i] = std::thread(single_thread_test_func[test_id], buf + offset, size, count / thread_num);
        } else { // random 
            test_threads[i] = std::thread(single_thread_test_func[test_id], buf, size, count / thread_num);
        }
    }
    // wait for all threads finishes
    for(int i = 0; i < thread_num; i++) {
        test_threads[i].join();
    }

    auto end = seconds();
    
    return end - start;
}

void print_help(char * name) {
    cout << "Usage: " << name <<" options" << endl;
    cout << "\t -p: Optane(0) or DRAM (1), default 0" << endl;
    cout << "\t -i: Test ID :" << endl;
    cout << "\t               (SEQ_RO, SEQ_RW41, SEQ_RW21, SEQ_RW11, SEQ_RW12, SEQ_RW14, SEQ_WO)" << endl;
    cout << "\t               (RND_RO, RND_RW41, RND_RW21, RND_RW11, RND_RW12, RND_RW14, RND_WO)" << endl;
    cout << "\t -t: Thread number, default 1" << endl;
    cout << "\t -g: Request unit, default 64" << endl;
    return ;
}

int main(int argc, char ** argv) {
    // default value for options
    bool opt_pmem = true;
    int opt_id = 0;
    int opt_threads = 1;
    int opt_granularity = 256;

    // read options from arguments
    static const char * optstr = "p:g:t:i:h";
    char opt;
    while((opt = getopt(argc, argv, optstr))!= -1) {
        switch(opt) {
            case 'p': opt_pmem = (atoi(optarg) == 0 ? true: false); break;
            case 'i': opt_id = atoi(optarg); break;
            case 't': opt_threads = atoi(optarg); break;
            case 'g': opt_granularity = (atoi(optarg) / 64) * 64; break;
            case 'h': 
            case '?': print_help(argv[0]); exit(-1);
            default : print_help(argv[0]); exit(-1);
        }
    }

    // print the options
    #ifdef DEBUG
        cout << "workload type: " << TEST_NAME[opt_id] << endl;
        cout << "thread number: " << opt_threads << endl;
        cout << "request unit : " << opt_granularity << "B" << endl;
        if(opt_pmem) {
            cout << "test environment: Optane" << endl;
        } else {
            cout << "test environment: DRAM" << endl;
        }
    #endif

    // preparing
    char * buf0 = NULL, * buf;
    size_t mapped_len;
    
    if (opt_pmem) {
        int is_pmem;
        if(file_exist("/mnt/pmem/bandwidth.pool")) {
            if(BUFSIZE > file_size("/mnt/pmem/bandwidth.pool")) {
                cout << "\nThe existing file is too small to do the test" << endl;
                exit(1);
            }
            // open the existing pool file and mmap it into virtual memory
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", 0, 0,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
        } else {
            // create a new pool file and mmap it into virtual memory 
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", BUFSIZE + 1024 * 1024, PMEM_FILE_CREATE,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
            
            // an empty file must be initialized with contents before reading
            omp_set_num_threads(16);
            #pragma omp parallel for
                for(int64_t i = 0; i < (BUFSIZE + 1024 * 1024) / 64; i++) {
                    buf0[i * 64] = 0x5a;
                }
        }

        // whether the pool file locates on PMEM or normal file system? 
        if(is_pmem == 0) {
            cout << "This file is not on pmem" << endl;
            exit(1);
        }

        buf = buf0;
    } else {
        buf0 = (char *) malloc(BUFSIZE + 1024 * 1024);
        uint64_t tmp = (uint64_t)buf0;
        uint64_t remaining = tmp & (ALIGN_SIZE - 1);
        buf = (char *) (remaining == 0 ? tmp : tmp + ALIGN_SIZE - remaining);
    }

    // the buffer address is already aligned 
    double elaspe_time = run_test(buf, opt_id, opt_granularity, BUFSIZE / opt_granularity, opt_threads);

    // output the result throughput
    //cout << elaspe_time << " ";
    cout << int(BUFSIZE / (1024 * 1024) / elaspe_time) << endl;

    // free memory
    if (opt_pmem)
        pmem_unmap(buf0, mapped_len);
    else
        free(buf0);
    
    return 0;
}
